# -*- coding: utf-8 -*-
"""
Created on Sat Dec 26 16:53:41 2015

@author: liuyi05
"""

import math
import pandas as pd

class Calc(object):
    LABEL_RETURN = "return"
    LABEL_ALLOT = "allot"
    LABEL_ATR = "atr"

    @staticmethod
    def calc_ma(df, period, label=None):
        assert(isinstance(period, int) and period >= 1)
        if label is None:
            label = "ma" + str(period)
        if label in df:
            return False
        close = df["close"]
        #ma = pd.rolling_mean(close, period)
        ma = close.rolling(period).mean()
#        for i in range(min(period - 1, df.shape[0])):
#            ma[i] = close[:i+1].mean()
        df[label] = ma
        return True

    @staticmethod
    def calc_return(df):
        if Calc.LABEL_RETURN in df and Calc.LABEL_ALLOT in df:
            return False
        close = df["close"]
        adj = df["adj_close"]
        cls_ret = close.pct_change()
        cls_ret[0] = 0
        adj_ret = adj.pct_change()
        adj_ret[0] = 0
        expect_cls = (close - close * adj_ret).shift(-1)
        allot = (close / expect_cls - 1).shift()
        allot[0] = 0
        df[Calc.LABEL_RETURN] = cls_ret
        df[Calc.LABEL_ALLOT] = 0
        delta = 0.02
        select = allot > delta
        df.ix[select, Calc.LABEL_RETURN] = adj_ret
        df.ix[select, Calc.LABEL_ALLOT] = allot
        return True

    @staticmethod
    def calc_atr(df):
        if Calc.LABEL_ATR in df:
            return False
        last_cls = df["close"].shift()
        tmp = pd.DataFrame({"last": last_cls, "high": df["high"], "low": df["low"]})
        atr = tmp[["last", "high"]].max(axis=1) - tmp[["low", "last"]].min(axis=1)
        df[Calc.LABEL_ATR] = atr
        return True

    @staticmethod
    def calc_atrxd(df, period, label=None):
        assert(isinstance(period, int) and period >= 1)
        if label is None:
            label = "atr" + str(period)
        if label in df:
            return False
        if Calc.LABEL_ATR not in df:
            Calc.calc_atr(df)
        if period == 1:
            return
        atr = df[Calc.LABEL_ATR]
        #atrxd = pd.rolling_mean(atr, period)
        atrxd = atr.rolling(period).mean()
#        for i in range(min(period - 1, df.shape[0])):
#            atrxd[i] = atr[:i+1].mean()
        df[label] = atrxd
        return True

    @staticmethod
    def calc_min(df, period, label=None):
        assert(isinstance(period, int) and period > 1)
        if label is None:
            label = "min" + str(period)
        if label in df:
            return False
        low = df["low"]
        #mm = pd.rolling_min(low, period).shift()
        mm = low.rolling(period).min().shift()
#        mm[0] = low[0]
#        for i in range(1, min(period, df.shape[0])):
#            mm[i] = low[:i].min()
        df[label] = mm
        return True

    @staticmethod
    def calc_max(df, period, label=None):
        assert(isinstance(period, int) and period > 1)
        if label is None:
            label = "max" + str(period)
        if label in df:
            return False
        high = df["high"]
        #mm = pd.rolling_max(high, period).shift()
        mm = high.rolling(period).max().shift()
#        mm[0] = high[0]
#        for i in range(1, min(period, df.shape[0])):
#            mm[i] = high[:i].max()
        df[label] = mm
        return True
        
    @staticmethod
    def calc_log_e_ratio(df, period, label=None):
        assert(isinstance(period, int) and period > 1)
        if label is None:
            label = "log_e_ratio" + str(period)
        if label in df:
            return False
        buy = df["close"]
        #highest = pd.rolling_max(df["high"], period).shift(-period)
        #lowest = pd.rolling_min(df["low"], period).shift(-period)
        highest = df["high"].rolling(period).max().shift(-period)
        lowest = df["low"].rolling(period).min().shift(-period)
        eps = 1e-6
        e_ratio = (pd.DataFrame([highest, buy]).max() - buy) \
                / (buy - pd.DataFrame([buy, lowest]).min() + eps)
        df[label] = e_ratio.apply(lambda x: math.log(x + 1))
        return True

    @staticmethod
    def calc_profit(df, period, label=None):
        assert(isinstance(period, int) and period > 1)
        if label is None:
            label = "prf" + str(period)
        if label in df:
            return False
        delta = 0.15
        open_price = df["open"]
        loss_limit = open_price * (1 - delta)
        #highest = pd.rolling_max(df["high"], period).shift(1-period)
        highest = df["high"].rolling(period).max().shift(1-period)
        prf = (highest - open_price) / open_price
        close = df["close"]
        for i in range(df.shape[0] - period + 1):
            for j in range(period):
                if close[i+j] <= loss_limit[i]:
                    prf[i] = (close[i+j] - open_price[i]) / open_price[i]
                    break
        df[label] = prf
        return True


def main():
    code = "000100"
    dir_path = "data/history/040101_051130/"
    path = dir_path + code + ".asc"
    df = pd.read_csv(path, index_col="date")

    Calc.calc_ma(df, 10, "ma10")
    Calc.calc_return(df)
    Calc.calc_atr(df)
    Calc.calc_atrxd(df, 10, "atr10")
    Calc.calc_min(df, 10, "min10")
    Calc.calc_max(df, 10, "max10")
    Calc.calc_log_e_ratio(df, 10, "log_e_ratio10")
    Calc.calc_profit(df, 10, "prf10")
    df.to_csv("haha.csv", float_format="%.2f")

if __name__ == "__main__":
    main()